import math
from einops import pack, unpack
import torch
import torch.nn.functional as F
from collections import namedtuple

# normalization functions

ModelPrediction =  namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])

def normalize_to_neg_one_to_one(img):
    return img * 2 - 1

def unnormalize_to_zero_to_one(t):
    return (t + 1) * 0.5

def exists(x):
    return x is not None

def identity(t, *args, **kwargs):
    return t

def default(val, d):
    '''
    if val not None, return it, then if d is callable, call it else return d
    '''
    if exists(val):
        return val
    return d() if callable(d) else d

def identity(t, *args, **kwargs):
    return t

def cycle(dl):
    '''
    cycle a data loader
    '''
    while True:
        for data in dl:
            yield data
            
def has_int_squareroot(num):
    return (math.sqrt(num)**2) == num

def num_to_groups(num, divisor):
    groups = num // divisor
    remainder = num % divisor
    arr = [divisor] * groups
    if remainder > 0:
        arr.append(remainder)
    return arr

def convert_image_to_fn(img_type, image):
    if image.mode != img_type:
        return image.convert(img_type)
    return image

def pack_one_with_inverse(x, pattern):
    packed, packed_shape = pack([x], pattern)

    def inverse(x, inverse_pattern = None):
        inverse_pattern = default(inverse_pattern, pattern)
        return unpack(x, packed_shape, inverse_pattern)[0]

    return packed, inverse

# image normalization functions
def normalize_to_neg_one_to_one(img):
    return img * 2 - 1

def unnormalize_to_zero_to_one(t):
    return (t+1) * 0.5

# classifier free guidance functions
def uniform(shape, device):
    return torch.zeros(shape, device=device).float().uniform_(0, 1)

def prob_mask_like(shape, prob, device):
    if prob == 1:
        return torch.ones(shape, device=device, dtype=torch.bool)
    elif prob == 0:
        return torch.zeros(shape, device=device, dtype=torch.bool)
    else:
        return torch.zeros(shape, device=device).float().uniform_(0,1) < prob
    
def project(x, y):
    x, inverse = pack_one_with_inverse(x, 'b *')
    y, _ = pack_one_with_inverse(y, 'b *')
    dtype = x.dtype
    x, y = x.double(), y.double()
    unit = F.normalize(y, dim=-1)
    parallel = (x * unit).sum(dim=-1, keepdim=True) * unit
    orthogonal = x - parallel
    
    return inverse(parallel).to(dtype), inverse(orthogonal).to(dtype)

def cast_tuple(t, length = 1):
    if isinstance(t, tuple):
        return t
    return ((t,) * length)

def divisible_by(numer, denom):
    return (numer % denom) == 0

def extract(a, t, x_shape):
    b, *_ = t.shape
    out = a.gather(-1, t)
    return out.reshape(b, *((1,) * (len(x_shape) - 1)))